package worker

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/fleetdm/fleet/v4/pkg/optjson"
	"github.com/fleetdm/fleet/v4/server/datastore/mysql"
	"github.com/fleetdm/fleet/v4/server/fleet"
	apple_mdm "github.com/fleetdm/fleet/v4/server/mdm/apple"
	nanomdm_push "github.com/fleetdm/fleet/v4/server/mdm/nanomdm/push"
	mock "github.com/fleetdm/fleet/v4/server/mock/mdm"
	"github.com/fleetdm/fleet/v4/server/ptr"
	"github.com/fleetdm/fleet/v4/server/test"
	kitlog "github.com/go-kit/log"
	"github.com/google/uuid"
	"github.com/jmoiron/sqlx"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type mockPusher struct {
	response *nanomdm_push.Response
	err      error
}

func (m mockPusher) Push(context.Context, []string) (map[string]*nanomdm_push.Response, error) {
	var res map[string]*nanomdm_push.Response
	if m.response != nil {
		res = map[string]*nanomdm_push.Response{
			m.response.Id: m.response,
		}
	}
	return res, m.err
}

type installAppResponse struct {
	CommandUUID string
	Error       error
}

type mockVPPInstaller struct {
	t                   *testing.T
	ds                  *mysql.Datastore
	installedApps       []*fleet.VPPApp
	appInstallResponses map[string]installAppResponse
	getTokenErr         error
}

func (m *mockVPPInstaller) GetVPPTokenIfCanInstallVPPApps(ctx context.Context, appleDevice bool, host *fleet.Host) (string, error) {
	require.True(m.t, appleDevice)
	if m.getTokenErr != nil {
		return "", m.getTokenErr
	}
	return "valid-token", nil
}

func (m *mockVPPInstaller) InstallVPPAppPostValidation(ctx context.Context, host *fleet.Host, vppApp *fleet.VPPApp, token string, opts fleet.HostSoftwareInstallOptions) (string, error) {
	require.True(m.t, opts.ForSetupExperience)
	resp, ok := m.appInstallResponses[vppApp.AdamID]
	require.True(m.t, ok)
	m.installedApps = append(m.installedApps, vppApp)
	if resp.Error == nil {
		mysql.ExecAdhocSQL(m.t, m.ds, func(q sqlx.ExtContext) error {
			_, err := q.ExecContext(ctx, `
          INSERT INTO nano_commands (command_uuid, request_type, command)
          VALUES (?, 'InstallApplication', '<?xml')
		`, resp.CommandUUID)
			if err != nil {
				return err
			}
			_, err = q.ExecContext(ctx, `
          INSERT INTO nano_enrollment_queue (id, command_uuid, active)
          VALUES (?, ?, 1)
		`, host.UUID, resp.CommandUUID)
			return err
		})
	}
	return resp.CommandUUID, resp.Error
}

func TestAppleMDM(t *testing.T) {
	ctx := context.Background()

	// use a real mysql datastore so that the test does not rely so much on
	// specific internals (sequence and number of calls, etc.). The MDM storage
	// and pusher are mocks.
	ds := mysql.CreateMySQLDS(t)
	// call TruncateTables immediately as a DB migation may have created jobs
	mysql.TruncateTables(t, ds)

	mdmStorage, err := ds.NewMDMAppleMDMStorage()
	require.NoError(t, err)

	// nopLog := kitlog.NewNopLogger()
	// use this to debug/verify details of calls
	nopLog := kitlog.NewJSONLogger(os.Stdout)

	testOrgName := "fleet-test"

	createEnrolledHost := func(t *testing.T, i int, teamID *uint, depAssignedToFleet bool, platform string) *fleet.Host {
		// create the host
		h, err := ds.NewHost(ctx, &fleet.Host{
			Hostname:       fmt.Sprintf("test-host%d-name", i),
			OsqueryHostID:  ptr.String(fmt.Sprintf("osquery-%d", i)),
			NodeKey:        ptr.String(fmt.Sprintf("nodekey-%d", i)),
			UUID:           uuid.New().String(),
			Platform:       platform,
			HardwareSerial: fmt.Sprintf("serial-%d", i),
			TeamID:         teamID,
		})
		require.NoError(t, err)

		// create the nano_device and enrollment
		var abmTokenID uint
		mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
			_, err := q.ExecContext(ctx, `INSERT INTO nano_devices (id, serial_number, authenticate) VALUES (?, ?, ?)`, h.UUID, h.HardwareSerial, "test")
			if err != nil {
				return err
			}
			_, err = q.ExecContext(ctx, `INSERT INTO nano_enrollments (id, device_id, type, topic, push_magic, token_hex, last_seen_at)
				VALUES (?, ?, ?, ?, ?, ?, ?)`, h.UUID, h.UUID, "device", "topic", "push_magic", "token_hex", time.Now())
			if err != nil {
				return err
			}

			encTok := uuid.NewString()
			abmToken, err := ds.InsertABMToken(ctx, &fleet.ABMToken{
				OrganizationName: "unused",
				EncryptedToken:   []byte(encTok),
				RenewAt:          time.Now().Add(30 * 24 * time.Hour), // 30 days from now
			})
			abmTokenID = abmToken.ID

			return err
		})
		if depAssignedToFleet {
			err := ds.UpsertMDMAppleHostDEPAssignments(ctx, []fleet.Host{*h}, abmTokenID, make(map[uint]time.Time))
			require.NoError(t, err)
		}
		err = ds.SetOrUpdateMDMData(ctx, h.ID, false, true, "http://example.com", depAssignedToFleet, fleet.WellKnownMDMFleet, "", false)
		require.NoError(t, err)
		return h
	}

	getEnqueuedCommandTypes := func(t *testing.T) []string {
		var commands []string
		mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
			return sqlx.SelectContext(ctx, q, &commands, "SELECT request_type FROM nano_commands")
		})
		return commands
	}

	enableManualRelease := func(t *testing.T, teamID *uint) {
		if teamID == nil {
			enableAppCfg := func(enable bool) {
				ac, err := ds.AppConfig(ctx)
				require.NoError(t, err)
				ac.MDM.MacOSSetup.EnableReleaseDeviceManually = optjson.SetBool(enable)
				err = ds.SaveAppConfig(ctx, ac)
				require.NoError(t, err)
			}

			enableAppCfg(true)
			t.Cleanup(func() { enableAppCfg(false) })
		} else {
			enableTm := func(enable bool) {
				tm, err := ds.TeamWithExtras(ctx, *teamID) // TODO see if we can convert to TeamLite (will require a new save DS method)
				require.NoError(t, err)
				tm.Config.MDM.MacOSSetup.EnableReleaseDeviceManually = optjson.SetBool(enable)
				_, err = ds.SaveTeam(ctx, tm)
				require.NoError(t, err)
			}

			enableTm(true)
			t.Cleanup(func() { enableTm(false) })
		}
	}

	t.Run("no-op with nil commander", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		// create a host and enqueue the job
		h := createEnrolledHost(t, 1, nil, true, "darwin")
		err := QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", nil, "", false)
		require.NoError(t, err)

		// run the worker, should mark the job as done
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)
		require.Empty(t, jobs)
	})

	t.Run("fails with unknown task", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		// create a host and enqueue the job
		h := createEnrolledHost(t, 1, nil, true, "darwin")
		err := QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMTask("no-such-task"), h.UUID, "darwin", nil, "", false)
		require.NoError(t, err)

		// run the worker, should mark the job as failed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Time{})
		require.NoError(t, err)
		require.Len(t, jobs, 1)
		require.Contains(t, jobs[0].Error, "unknown task: no-such-task")
		require.Equal(t, fleet.JobStateQueued, jobs[0].State)
		require.Equal(t, 1, jobs[0].Retries)
	})

	t.Run("installs default manifest", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		h := createEnrolledHost(t, 1, nil, true, "darwin")

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		// use "" instead of "darwin" as platform to test a queued job after the upgrade to iOS/iPadOS support.
		err := QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "", nil, "", false)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)

		// there is no post-DEP release device job anymore
		require.Len(t, jobs, 0)

		require.ElementsMatch(t, []string{"InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))
	})

	t.Run("installs default manifest, manual release", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		t.Cleanup(func() { mysql.TruncateTables(t, ds) })

		h := createEnrolledHost(t, 1, nil, true, "darwin")
		enableManualRelease(t, nil)

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err = QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", nil, "", false)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)

		// there is no post-DEP release device job pending
		require.Empty(t, jobs)

		require.ElementsMatch(t, []string{"InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))
	})

	t.Run("installs custom bootstrap manifest", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		h := createEnrolledHost(t, 1, nil, true, "darwin")
		err := ds.InsertMDMAppleBootstrapPackage(ctx, &fleet.MDMAppleBootstrapPackage{
			Name:   "custom-bootstrap",
			TeamID: 0, // no-team
			Bytes:  []byte("test"),
			Sha256: []byte("test"),
			Token:  "token",
		}, nil)
		require.NoError(t, err)

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err = QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", nil, "", false)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)

		// the post-DEP release device job is not queued anymore
		require.Len(t, jobs, 0)

		require.ElementsMatch(t, []string{"InstallEnterpriseApplication", "InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))

		ms, err := ds.GetHostMDMMacOSSetup(ctx, h.ID)
		require.NoError(t, err)
		require.Equal(t, "custom-bootstrap", ms.BootstrapPackageName)
	})

	t.Run("installs custom bootstrap manifest of a team", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		tm, err := ds.NewTeam(ctx, &fleet.Team{Name: "test"})
		require.NoError(t, err)

		h := createEnrolledHost(t, 1, &tm.ID, true, "darwin")
		err = ds.InsertMDMAppleBootstrapPackage(ctx, &fleet.MDMAppleBootstrapPackage{
			Name:   "custom-team-bootstrap",
			TeamID: tm.ID,
			Bytes:  []byte("test"),
			Sha256: []byte("test"),
			Token:  "token",
		}, nil)
		require.NoError(t, err)

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err = QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", &tm.ID, "", false)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)

		// the post-DEP release device job is not queued anymore
		require.Len(t, jobs, 0)

		require.ElementsMatch(t, []string{"InstallEnterpriseApplication", "InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))

		ms, err := ds.GetHostMDMMacOSSetup(ctx, h.ID)
		require.NoError(t, err)
		require.Equal(t, "custom-team-bootstrap", ms.BootstrapPackageName)
	})

	t.Run("installs custom bootstrap manifest of a team, manual release", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		t.Cleanup(func() { mysql.TruncateTables(t, ds) })

		tm, err := ds.NewTeam(ctx, &fleet.Team{Name: "test"})
		require.NoError(t, err)
		enableManualRelease(t, &tm.ID)

		h := createEnrolledHost(t, 1, &tm.ID, true, "darwin")
		err = ds.InsertMDMAppleBootstrapPackage(ctx, &fleet.MDMAppleBootstrapPackage{
			Name:   "custom-team-bootstrap",
			TeamID: tm.ID,
			Bytes:  []byte("test"),
			Sha256: []byte("test"),
			Token:  "token",
		}, nil)
		require.NoError(t, err)

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err = QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", &tm.ID, "", false)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)

		// there is no post-DEP release device job pending
		require.Empty(t, jobs)

		require.ElementsMatch(t, []string{"InstallEnterpriseApplication", "InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))

		ms, err := ds.GetHostMDMMacOSSetup(ctx, h.ID)
		require.NoError(t, err)
		require.Equal(t, "custom-team-bootstrap", ms.BootstrapPackageName)
	})

	t.Run("unknown enroll reference", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		h := createEnrolledHost(t, 1, nil, true, "darwin")

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err := QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", nil, "abcd", false)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Time{})
		require.NoError(t, err)
		require.Len(t, jobs, 1)
		require.Contains(t, jobs[0].Error, "MDMIdPAccount with uuid abcd was not found")
		require.Equal(t, fleet.JobStateQueued, jobs[0].State)
		require.Equal(t, 1, jobs[0].Retries)
	})

	t.Run("enroll reference but SSO disabled", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		err := ds.InsertMDMIdPAccount(ctx, &fleet.MDMIdPAccount{
			Username: "test",
			Fullname: "test",
			Email:    "test@example.com",
		})
		require.NoError(t, err)

		idpAcc, err := ds.GetMDMIdPAccountByEmail(ctx, "test@example.com")
		require.NoError(t, err)

		h := createEnrolledHost(t, 1, nil, true, "darwin")

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err = QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", nil, idpAcc.UUID, false)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)

		// the post-DEP release device job is not queued anymore
		require.Len(t, jobs, 0)

		// confirm that AccountConfiguration command was not enqueued
		require.ElementsMatch(t, []string{"InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))
	})

	t.Run("enroll reference with SSO enabled", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		err := ds.InsertMDMIdPAccount(ctx, &fleet.MDMIdPAccount{
			Username: "test",
			Fullname: "test",
			Email:    "test@example.com",
		})
		require.NoError(t, err)

		idpAcc, err := ds.GetMDMIdPAccountByEmail(ctx, "test@example.com")
		require.NoError(t, err)

		tm, err := ds.NewTeam(ctx, &fleet.Team{Name: "test"})
		require.NoError(t, err)
		tm, err = ds.TeamWithExtras(ctx, tm.ID) // TODO see if we can convert to TeamLite (will require a new save DS method)
		require.NoError(t, err)
		tm.Config.MDM.MacOSSetup.EnableEndUserAuthentication = true
		_, err = ds.SaveTeam(ctx, tm)
		require.NoError(t, err)

		h := createEnrolledHost(t, 1, &tm.ID, true, "darwin")

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err = QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", &tm.ID, idpAcc.UUID, false)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)

		// the post-DEP release device job is not queued anymore
		require.Len(t, jobs, 0)

		require.ElementsMatch(t, []string{"InstallEnterpriseApplication", "AccountConfiguration"}, getEnqueuedCommandTypes(t))
	})

	t.Run("installs fleetd for manual enrollments", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		h := createEnrolledHost(t, 1, nil, true, "darwin")

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err := QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostManualEnrollmentTask, h.UUID, "darwin", nil, "", false)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)
		require.Empty(t, jobs)
		require.ElementsMatch(t, []string{"InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))
	})

	t.Run("use worker for automatic release", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		h := createEnrolledHost(t, 1, nil, true, "darwin")

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err := QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", nil, "", true)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		require.ElementsMatch(t, []string{"InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))

		// the release device job got enqueued
		jobs, err := ds.GetQueuedJobs(ctx, 1, time.Now().Add(time.Minute)) // release job is always added with a delay
		require.NoError(t, err)
		require.Len(t, jobs, 1)
		require.Equal(t, fleet.JobStateQueued, jobs[0].State)
		require.Equal(t, appleMDMJobName, jobs[0].Name)
		require.Contains(t, string(*jobs[0].Args), AppleMDMPostDEPReleaseDeviceTask)
	})

	t.Run("automatic release retries and give up", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		h := createEnrolledHost(t, 1, nil, true, "darwin")

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err := QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", nil, "", true)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		require.ElementsMatch(t, []string{"InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))

		// the release device job got enqueued, and it will constantly re-enqueue
		// itself because the command is never acknowledged
		var (
			previousID     uint
			firstStartedAt time.Time
		)
		for i := 0; i <= 10; i++ {
			jobs, err := ds.GetQueuedJobs(ctx, 2, time.Now().UTC().Add(time.Minute)) // release job is always added with a delay
			require.NoError(t, err)
			require.Len(t, jobs, 1)

			releaseJob := jobs[0]
			require.Equal(t, fleet.JobStateQueued, releaseJob.State)
			require.Equal(t, appleMDMJobName, releaseJob.Name)
			require.NotEqual(t, previousID, releaseJob.ID)
			previousID = releaseJob.ID

			var args appleMDMArgs
			err = json.Unmarshal([]byte(*releaseJob.Args), &args)
			require.NoError(t, err)
			require.Equal(t, args.Task, AppleMDMPostDEPReleaseDeviceTask)
			require.EqualValues(t, i, args.ReleaseDeviceAttempt)

			if i == 0 {
				// first time, there is no release device started at
				require.Nil(t, args.ReleaseDeviceStartedAt)
			} else {
				require.NotNil(t, args.ReleaseDeviceStartedAt)
				if i == 1 {
					firstStartedAt = *args.ReleaseDeviceStartedAt
				} else {
					require.True(t, firstStartedAt.Equal(*args.ReleaseDeviceStartedAt))
				}
			}

			if i == 10 {
				// finally, after 10 attempts, update the release started at to make it
				// meet the maximum wait time and actually do the release on the next
				// processing.
				startedAt := firstStartedAt.Add(-time.Hour)
				args.ReleaseDeviceStartedAt = &startedAt
				b, err := json.Marshal(args)
				require.NoError(t, err)
				mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
					_, err := q.ExecContext(ctx, `UPDATE jobs SET args = ? WHERE id = ?`, string(b), releaseJob.ID)
					return err
				})
			}
			// update the job to make it available to run immediately
			releaseJob.NotBefore = time.Now().UTC().Add(-time.Minute)
			_, err = ds.UpdateJob(ctx, releaseJob.ID, releaseJob)
			require.NoError(t, err)

			// run the worker, should succeed and re-enqueue a new job with the same args
			err = w.ProcessJobs(ctx)
			require.NoError(t, err)
		}

		// on the last processing, it did end up releasing the device due to the
		// limit of attempts and wait delay being reached.
		require.ElementsMatch(t, []string{"InstallEnterpriseApplication", "DeviceConfigured"}, getEnqueuedCommandTypes(t))

		// job queue is now empty
		jobs, err := ds.GetQueuedJobs(ctx, 2, time.Now().UTC().Add(time.Minute))
		require.NoError(t, err)
		require.Len(t, jobs, 0)
	})

	t.Run("automatic release succeeds after a few attempts", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		defer mysql.TruncateTables(t, ds)

		h := createEnrolledHost(t, 1, nil, true, "darwin")

		mdmWorker := &AppleMDM{
			Datastore: ds,
			Log:       nopLog,
			Commander: apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err := QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, "darwin", nil, "", true)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		require.ElementsMatch(t, []string{"InstallEnterpriseApplication"}, getEnqueuedCommandTypes(t))

		for i := 0; i <= 4; i++ {
			jobs, err := ds.GetQueuedJobs(ctx, 2, time.Now().UTC().Add(time.Minute)) // release job is always added with a delay
			require.NoError(t, err)
			require.Len(t, jobs, 1)

			releaseJob := jobs[0]
			require.Equal(t, fleet.JobStateQueued, releaseJob.State)
			require.Equal(t, appleMDMJobName, releaseJob.Name)

			if i == 4 {
				// after 4 attempts, record a result for the command so it gets released
				mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
					_, err := q.ExecContext(ctx, `INSERT INTO nano_command_results (id, command_uuid, status, result) 
						SELECT ?, command_uuid, ?, ? FROM nano_commands`,
						h.UUID, "Acknowledged", `<?xml`)
					return err
				})
			}
			// update the job to make it available to run immediately
			releaseJob.NotBefore = time.Now().UTC().Add(-time.Minute)
			_, err = ds.UpdateJob(ctx, releaseJob.ID, releaseJob)
			require.NoError(t, err)

			// run the worker, should succeed and re-enqueue a new job with the same args
			err = w.ProcessJobs(ctx)
			require.NoError(t, err)
		}

		// on the last processing, it did release the device due to all pending
		// commands being completed.
		require.ElementsMatch(t, []string{"InstallEnterpriseApplication", "DeviceConfigured"}, getEnqueuedCommandTypes(t))

		// job queue is now empty
		jobs, err := ds.GetQueuedJobs(ctx, 2, time.Now().UTC().Add(time.Minute))
		require.NoError(t, err)
		require.Len(t, jobs, 0)
	})

	t.Run("installs enqueued VPP apps", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		test.CreateInsertGlobalVPPToken(t, ds)
		defer mysql.TruncateTables(t, ds)

		tm, err := ds.NewTeam(ctx, &fleet.Team{Name: "test"})
		require.NoError(t, err)

		h := createEnrolledHost(t, 1, &tm.ID, true, "ios")

		expectedAppInstalls := []*fleet.VPPApp{}

		for i := 0; i < 3; i++ {
			idx := fmt.Sprint(i)
			vppApp := &fleet.VPPApp{
				Name: "vpp_worker-" + idx, LatestVersion: "1.0.0", VPPAppTeam: fleet.VPPAppTeam{VPPAppID: fleet.VPPAppID{AdamID: "depworker-" + idx, Platform: fleet.IOSPlatform}},
				BundleIdentifier: "b" + idx,
			}
			vppAppWithTeam, err := ds.InsertVPPAppWithTeam(ctx, vppApp, &tm.ID)
			require.NoError(t, err)
			expectedAppInstalls = append(expectedAppInstalls, vppAppWithTeam)
		}

		appInstallResponses := make(map[string]installAppResponse, len(expectedAppInstalls))

		for _, appWithTeam := range expectedAppInstalls {
			mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
				stmt := `
INSERT INTO setup_experience_status_results (
	host_uuid,
	name,
	status,
	vpp_app_team_id
) VALUES (?, ?, ?, ?)
 `
				_, err = q.ExecContext(ctx, stmt, h.UUID, appWithTeam.Name, fleet.SetupExperienceStatusPending, appWithTeam.VPPAppTeam.AppTeamID)
				return err
			})
			appInstallResponses[appWithTeam.AdamID] = installAppResponse{CommandUUID: uuid.NewString(), Error: nil}
		}

		vppInstaller := &mockVPPInstaller{t: t, ds: ds, appInstallResponses: appInstallResponses}

		mdmWorker := &AppleMDM{
			VPPInstaller: vppInstaller,
			Datastore:    ds,
			Log:          nopLog,
			Commander:    apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err = QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, h.Platform, nil, "", true)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 10, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)
		require.NotEmpty(t, jobs)
		var releaseJob *fleet.Job
		for _, job := range jobs {
			if job.Name == appleMDMJobName {
				// THere should only be one release job
				require.Nil(t, releaseJob)
				releaseJob = job
			}
		}
		// We should have found a release job
		require.NotNil(t, releaseJob)
		// It should be the release task
		require.Contains(t, string(*releaseJob.Args), AppleMDMPostDEPReleaseDeviceTask)
		// And it should contain the command IDs for the installs
		expectedAdamIDs := make([]string, 0, len(expectedAppInstalls))
		installedAdamIDs := make([]string, 0, len(vppInstaller.installedApps))
		for _, app := range expectedAppInstalls {
			require.Contains(t, string(*releaseJob.Args), appInstallResponses[app.AdamID].CommandUUID)
			expectedAdamIDs = append(expectedAdamIDs, app.AdamID)
		}
		for _, installed := range vppInstaller.installedApps {
			installedAdamIDs = append(installedAdamIDs, installed.AdamID)
		}
		require.ElementsMatch(t, expectedAdamIDs, installedAdamIDs)

		results, err := ds.ListSetupExperienceResultsByHostUUID(ctx, h.UUID)
		require.NoError(t, err)
		require.Len(t, results, len(expectedAppInstalls))
		for _, result := range results {
			require.Equal(t, fleet.SetupExperienceStatusRunning, result.Status)
		}

		// Acknowledge the commands - the release job should still re-enqueue itself and await the installs
		mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
			_, err := q.ExecContext(ctx, `INSERT INTO nano_command_results (id, command_uuid, status, result) 
						SELECT ?, command_uuid, ?, ? FROM nano_commands`,
				h.UUID, "Acknowledged", `<?xml`)
			return err
		})

		mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
			_, err := q.ExecContext(ctx, `UPDATE jobs SET not_before=? WHERE id=?`, time.Now().Add(-time.Minute), releaseJob.ID)
			return err
		})

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		releaseJob = nil

		jobs, err = ds.GetQueuedJobs(ctx, 10, time.Now().UTC().Add(time.Minute+time.Second)) // look in the future to catch any delayed job

		for _, job := range jobs {
			if job.Name == appleMDMJobName {
				// THere should only be one release job
				require.Nil(t, releaseJob)
				releaseJob = job
			}
		}
		// We should have found a release job
		require.NotNil(t, releaseJob)
		// It should be the release task
		require.Contains(t, string(*releaseJob.Args), AppleMDMPostDEPReleaseDeviceTask)

		// Now update setup_experience_status as if the installs succeeded
		mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
			_, err := q.ExecContext(ctx, `UPDATE setup_experience_status_results SET status=? WHERE host_uuid=?`, fleet.SetupExperienceStatusSuccess, h.UUID)
			return err
		})

		mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
			_, err := q.ExecContext(ctx, `UPDATE jobs SET not_before=? WHERE id=?`, time.Now().Add(-time.Minute), releaseJob.ID)
			return err
		})

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		jobs, err = ds.GetQueuedJobs(ctx, 10, time.Now().UTC().Add(time.Minute+time.Second)) // look in the future to catch any delayed job

		for _, job := range jobs {
			if job.Name == appleMDMJobName {
				require.Fail(t, "there should be no more release jobs queued")
			}
		}

		require.Contains(t, getEnqueuedCommandTypes(t), "DeviceConfigured")
	})

	t.Run("marks failed VPP installs as failed, runs all others", func(t *testing.T) {
		mysql.SetTestABMAssets(t, ds, testOrgName)
		test.CreateInsertGlobalVPPToken(t, ds)
		defer mysql.TruncateTables(t, ds)
		badCommandUUID := "bad-command-uuid"

		tm, err := ds.NewTeam(ctx, &fleet.Team{Name: "test"})
		require.NoError(t, err)

		h := createEnrolledHost(t, 1, &tm.ID, true, "ios")

		expectedAppInstalls := []*fleet.VPPApp{}

		for i := 0; i < 3; i++ {
			idx := fmt.Sprint(i)
			vppApp := &fleet.VPPApp{
				Name: "vpp_worker-" + idx, LatestVersion: "1.0.0", VPPAppTeam: fleet.VPPAppTeam{VPPAppID: fleet.VPPAppID{AdamID: "depworker-" + idx, Platform: fleet.IOSPlatform}},
				BundleIdentifier: "b" + idx,
			}
			vppAppWithTeam, err := ds.InsertVPPAppWithTeam(ctx, vppApp, &tm.ID)
			require.NoError(t, err)
			expectedAppInstalls = append(expectedAppInstalls, vppAppWithTeam)
		}

		appInstallResponses := make(map[string]installAppResponse, len(expectedAppInstalls))

		for _, appWithTeam := range expectedAppInstalls {
			mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
				stmt := `
INSERT INTO setup_experience_status_results (
	host_uuid,
	name,
	status,
	vpp_app_team_id
) VALUES (?, ?, ?, ?)
 `
				_, err = q.ExecContext(ctx, stmt, h.UUID, appWithTeam.Name, fleet.SetupExperienceStatusPending, appWithTeam.VPPAppTeam.AppTeamID)
				return err
			})
			if len(appInstallResponses) == 0 {
				// first one, simulate a failure. It shouldn't actually
				// return a command UUID here but even if it does we
				// should not wait on it
				appInstallResponses[appWithTeam.AdamID] = installAppResponse{CommandUUID: badCommandUUID, Error: errors.New("test error")}
				continue
			}
			// rest succeed
			appInstallResponses[appWithTeam.AdamID] = installAppResponse{CommandUUID: uuid.NewString(), Error: nil}
		}

		vppInstaller := &mockVPPInstaller{t: t, ds: ds, appInstallResponses: appInstallResponses}

		mdmWorker := &AppleMDM{
			VPPInstaller: vppInstaller,
			Datastore:    ds,
			Log:          nopLog,
			Commander:    apple_mdm.NewMDMAppleCommander(mdmStorage, mockPusher{}),
		}
		w := NewWorker(ds, nopLog)
		w.Register(mdmWorker)

		err = QueueAppleMDMJob(ctx, ds, nopLog, AppleMDMPostDEPEnrollmentTask, h.UUID, h.Platform, nil, "", true)
		require.NoError(t, err)

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		// ensure the job's not_before allows it to be returned if it were to run
		// again
		time.Sleep(time.Second)

		jobs, err := ds.GetQueuedJobs(ctx, 10, time.Now().UTC().Add(time.Minute)) // look in the future to catch any delayed job
		require.NoError(t, err)
		require.NotEmpty(t, jobs)
		var releaseJob *fleet.Job
		for _, job := range jobs {
			if job.Name == appleMDMJobName {
				// THere should only be one release job
				require.Nil(t, releaseJob)
				releaseJob = job
			}
		}
		// We should have found a release job
		require.NotNil(t, releaseJob)
		// It should be the release task
		require.Contains(t, string(*releaseJob.Args), AppleMDMPostDEPReleaseDeviceTask)

		// And it should contain the command IDs for the installs that didn't error
		expectedAdamIDs := make([]string, 0, len(expectedAppInstalls))
		installedAdamIDs := make([]string, 0, len(vppInstaller.installedApps))
		for _, app := range expectedAppInstalls {
			expectedAdamIDs = append(expectedAdamIDs, app.AdamID)
			if appInstallResponses[app.AdamID].Error != nil {
				// this one failed, so it should not be in the release command
				continue
			}
			require.Contains(t, string(*releaseJob.Args), appInstallResponses[app.AdamID].CommandUUID)
		}
		require.NotContains(t, string(*releaseJob.Args), badCommandUUID)
		for _, installed := range vppInstaller.installedApps {
			installedAdamIDs = append(installedAdamIDs, installed.AdamID)
		}
		require.ElementsMatch(t, expectedAdamIDs, installedAdamIDs)

		results, err := ds.ListSetupExperienceResultsByHostUUID(ctx, h.UUID)
		require.NoError(t, err)
		require.Len(t, results, len(expectedAppInstalls))
		for _, result := range results {
			require.NotNil(t, result.VPPAppAdamID)
			if *result.VPPAppAdamID == expectedAppInstalls[0].AdamID {
				// this is the one we simulated a failure for
				require.Equal(t, fleet.SetupExperienceStatusFailure, result.Status)
				continue
			}
			require.Equal(t, fleet.SetupExperienceStatusRunning, result.Status)
		}

		// Acknowledge the commands - the release job should still re-enqueue itself and await the remaining installs
		mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
			_, err := q.ExecContext(ctx, `INSERT INTO nano_command_results (id, command_uuid, status, result) 
						SELECT ?, command_uuid, ?, ? FROM nano_commands`,
				h.UUID, "Acknowledged", `<?xml`)
			return err
		})

		mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
			_, err := q.ExecContext(ctx, `UPDATE jobs SET not_before=? WHERE id=?`, time.Now().Add(-time.Minute), releaseJob.ID)
			return err
		})

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		releaseJob = nil

		jobs, err = ds.GetQueuedJobs(ctx, 10, time.Now().UTC().Add(time.Minute+time.Second)) // look in the future to catch any delayed job

		for _, job := range jobs {
			if job.Name == appleMDMJobName {
				// THere should only be one release job
				require.Nil(t, releaseJob)
				releaseJob = job
			}
		}
		// We should have found a release job
		require.NotNil(t, releaseJob)
		// It should be the release task
		require.Contains(t, string(*releaseJob.Args), AppleMDMPostDEPReleaseDeviceTask)

		// Now update setup_experience_status as if the installs succeeded
		mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
			_, err := q.ExecContext(ctx, `UPDATE setup_experience_status_results SET status=? WHERE host_uuid=? AND status <> ?`, fleet.SetupExperienceStatusSuccess, h.UUID, fleet.SetupExperienceStatusFailure)
			return err
		})

		mysql.ExecAdhocSQL(t, ds, func(q sqlx.ExtContext) error {
			_, err := q.ExecContext(ctx, `UPDATE jobs SET not_before=? WHERE id=?`, time.Now().Add(-time.Minute), releaseJob.ID)
			return err
		})

		// run the worker, should succeed
		err = w.ProcessJobs(ctx)
		require.NoError(t, err)

		jobs, err = ds.GetQueuedJobs(ctx, 10, time.Now().UTC().Add(time.Minute+time.Second)) // look in the future to catch any delayed job

		for _, job := range jobs {
			if job.Name == appleMDMJobName {
				assert.Fail(t, "there should be no more release jobs queued")
			}
		}

		require.Contains(t, getEnqueuedCommandTypes(t), "DeviceConfigured")
	})
}

func TestGetSignedURL(t *testing.T) {
	t.Parallel()

	ctx := context.Background()
	meta := &fleet.MDMAppleBootstrapPackage{
		Sha256: []byte{1, 2, 3},
	}

	var data []byte
	buf := bytes.NewBuffer(data)
	logger := kitlog.NewLogfmtLogger(buf)
	a := &AppleMDM{Log: logger}

	// S3 not configured
	assert.Empty(t, a.getSignedURL(ctx, meta))
	assert.Empty(t, buf.String())

	// Signer not configured
	mockStore := &mock.MDMBootstrapPackageStore{}
	a.BootstrapPackageStore = mockStore
	mockStore.SignFunc = func(ctx context.Context, fileID string) (string, error) {
		return "bozo", fleet.ErrNotConfigured
	}
	assert.Empty(t, a.getSignedURL(ctx, meta))
	assert.Empty(t, buf.String())

	// Test happy path
	mockStore.SignFunc = func(ctx context.Context, fileID string) (string, error) {
		return "signed", nil
	}
	mockStore.ExistsFunc = func(ctx context.Context, packageID string) (bool, error) {
		assert.Equal(t, "010203", packageID)
		return true, nil
	}
	assert.Equal(t, "signed", a.getSignedURL(ctx, meta))
	assert.Empty(t, buf.String())
	assert.True(t, mockStore.SignFuncInvoked)
	assert.True(t, mockStore.ExistsFuncInvoked)
	mockStore.SignFuncInvoked = false
	mockStore.ExistsFuncInvoked = false

	// Test error -- sign failed
	mockStore.SignFunc = func(ctx context.Context, fileID string) (string, error) {
		return "", errors.New("test error")
	}
	assert.Empty(t, a.getSignedURL(ctx, meta))
	assert.Contains(t, buf.String(), "test error")
	assert.True(t, mockStore.SignFuncInvoked)
	assert.False(t, mockStore.ExistsFuncInvoked)
}
